{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Coding: LeNet for MNIST"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device: cpu\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 9.91M/9.91M [00:00<00:00, 14.6MB/s]\n",
      "100%|██████████| 28.9k/28.9k [00:00<00:00, 394kB/s]\n",
      "100%|██████████| 1.65M/1.65M [00:00<00:00, 3.69MB/s]\n",
      "100%|██████████| 4.54k/4.54k [00:00<00:00, 11.1MB/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting training...\n",
      "Epoch [1/10]:\n",
      "Train Loss: 0.2402, Train Acc: 92.56%\n",
      "Test Loss: 0.0785, Test Acc: 97.49%\n",
      "--------------------------------------------------\n",
      "Epoch [2/10]:\n",
      "Train Loss: 0.0705, Train Acc: 97.72%\n",
      "Test Loss: 0.0471, Test Acc: 98.44%\n",
      "--------------------------------------------------\n",
      "Epoch [3/10]:\n",
      "Train Loss: 0.0510, Train Acc: 98.40%\n",
      "Test Loss: 0.0490, Test Acc: 98.40%\n",
      "--------------------------------------------------\n",
      "Epoch [4/10]:\n",
      "Train Loss: 0.0383, Train Acc: 98.81%\n",
      "Test Loss: 0.0384, Test Acc: 98.78%\n",
      "--------------------------------------------------\n",
      "Epoch [5/10]:\n",
      "Train Loss: 0.0333, Train Acc: 98.95%\n",
      "Test Loss: 0.0388, Test Acc: 98.80%\n",
      "--------------------------------------------------\n",
      "Epoch [6/10]:\n",
      "Train Loss: 0.0289, Train Acc: 99.06%\n",
      "Test Loss: 0.0355, Test Acc: 98.85%\n",
      "--------------------------------------------------\n",
      "Epoch [7/10]:\n",
      "Train Loss: 0.0237, Train Acc: 99.26%\n",
      "Test Loss: 0.0306, Test Acc: 99.08%\n",
      "--------------------------------------------------\n",
      "Epoch [8/10]:\n",
      "Train Loss: 0.0206, Train Acc: 99.34%\n",
      "Test Loss: 0.0371, Test Acc: 98.93%\n",
      "--------------------------------------------------\n",
      "Epoch [9/10]:\n",
      "Train Loss: 0.0187, Train Acc: 99.38%\n",
      "Test Loss: 0.0455, Test Acc: 98.63%\n",
      "--------------------------------------------------\n",
      "Epoch [10/10]:\n",
      "Train Loss: 0.0160, Train Acc: 99.46%\n",
      "Test Loss: 0.0314, Test Acc: 99.15%\n",
      "--------------------------------------------------\n",
      "Training completed and model saved!\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "# Set random seed for reproducibility\n",
    "torch.manual_seed(42)\n",
    "\n",
    "# Define the LeNet architecture\n",
    "class LeNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(LeNet, self).__init__()\n",
    "        self.conv0 = nn.Conv2d(1, 6, kernel_size=5)\n",
    "        self.conv1 = nn.Conv2d(6, 16, kernel_size=5)\n",
    "        self.fc0 = nn.Linear(16 * 4 * 4, 120)\n",
    "        self.fc1 = nn.Linear(120, 84)\n",
    "        self.fc2 = nn.Linear(84, 10)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Single-Channel 1-6: First conv block\n",
    "        # dimension: 28x28x1 -> 24x24x6\n",
    "        x = self.relu(self.conv0(x))\n",
    "        # dimension: 24x24x6 -> 12x12x6\n",
    "        x = self.maxpool(x)\n",
    "        \n",
    "        # Muti-Channel 6-16: Second conv block\n",
    "        # dimension: 12x12x6 -> 8x8x16\n",
    "        x = self.relu(self.conv1(x))\n",
    "        # dimension: 8x8x16 -> 4x4x16\n",
    "        x = self.maxpool(x)\n",
    "        \n",
    "        # MLP: Flatten and fully connected layers\n",
    "        x = x.view(x.size(0), -1)\n",
    "        # dimension: 16x4x4 -> 120\n",
    "        x = self.relu(self.fc0(x))\n",
    "        # dimension: 120 -> 84\n",
    "        x = self.relu(self.fc1(x))\n",
    "        # dimension: 84 -> 10\n",
    "        x = self.fc2(x)\n",
    "        return x\n",
    "\n",
    "# Data loading and preprocessing\n",
    "def load_data(batch_size=64):\n",
    "    transform = transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.1307,), (0.3081,))\n",
    "    ])\n",
    "    \n",
    "    train_dataset = torchvision.datasets.MNIST(\n",
    "        root='./data', \n",
    "        train=True,\n",
    "        download=True, \n",
    "        transform=transform\n",
    "    )\n",
    "    \n",
    "    test_dataset = torchvision.datasets.MNIST(\n",
    "        root='./data', \n",
    "        train=False,\n",
    "        download=True, \n",
    "        transform=transform\n",
    "    )\n",
    "    \n",
    "    train_loader = DataLoader(\n",
    "        train_dataset, \n",
    "        batch_size=batch_size,\n",
    "        shuffle=True\n",
    "    )\n",
    "    \n",
    "    test_loader = DataLoader(\n",
    "        test_dataset, \n",
    "        batch_size=batch_size,\n",
    "        shuffle=False\n",
    "    )\n",
    "    \n",
    "    return train_loader, test_loader\n",
    "\n",
    "# Training\n",
    "def train(model, train_loader, criterion, optimizer, device):\n",
    "    model.train()\n",
    "    running_loss = 0.0\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    \n",
    "    for images, labels in train_loader:\n",
    "        images, labels = images.to(device), labels.to(device)\n",
    "        \n",
    "        # Do Not Forget to Zero Gradients\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(images)\n",
    "        loss = criterion(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        running_loss += loss.item()\n",
    "        _, predicted = outputs.max(1)\n",
    "        total += labels.size(0)\n",
    "        correct += predicted.eq(labels).sum().item()\n",
    "    \n",
    "    accuracy = 100. * correct / total\n",
    "    return running_loss / len(train_loader), accuracy\n",
    "\n",
    "# Evaluation\n",
    "def evaluate(model, test_loader, criterion, device):\n",
    "    model.eval()\n",
    "    running_loss = 0.0\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for images, labels in test_loader:\n",
    "            images, labels = images.to(device), labels.to(device)\n",
    "            outputs = model(images)\n",
    "            loss = criterion(outputs, labels)\n",
    "            \n",
    "            running_loss += loss.item()\n",
    "            _, predicted = outputs.max(1)\n",
    "            total += labels.size(0)\n",
    "            correct += predicted.eq(labels).sum().item()\n",
    "    \n",
    "    accuracy = 100. * correct / total\n",
    "    return running_loss / len(test_loader), accuracy\n",
    "\n",
    "def main():\n",
    "    # Set device\n",
    "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "    print(f'Using device: {device}')\n",
    "    \n",
    "    # Hyperparameters\n",
    "    batch_size = 64\n",
    "    learning_rate = 0.001\n",
    "    num_epochs = 10\n",
    "    \n",
    "    # Load data\n",
    "    train_loader, test_loader = load_data(batch_size)\n",
    "    \n",
    "    # Initialize model, loss function, and optimizer\n",
    "    model = LeNet().to(device)\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n",
    "    \n",
    "    # Training loop\n",
    "    print('Starting training...')\n",
    "    for epoch in range(num_epochs):\n",
    "        train_loss, train_acc = train(model, train_loader, criterion, optimizer, device)\n",
    "        test_loss, test_acc = evaluate(model, test_loader, criterion, device)\n",
    "        \n",
    "        print(f'Epoch [{epoch+1}/{num_epochs}]:')\n",
    "        print(f'Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}%')\n",
    "        print(f'Test Loss: {test_loss:.4f}, Test Acc: {test_acc:.2f}%')\n",
    "        print('-' * 50)\n",
    "    \n",
    "    # Save the trained model\n",
    "    torch.save(model.state_dict(), 'lenet_mnist.pth')\n",
    "    print('Training completed and model saved!')\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "archer",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
